#pragma warning(disable : 4355)
#include "Avalon/Network/MulticastSocket.hpp"
#include <boost/asio.hpp>
#include <boost/bind.hpp>
#include "Avalon/IO/Buffer.hpp"
#include "Avalon/Network/SocketException.hpp"
#include "Avalon/Network/SocketThreadPool.hpp"

using namespace Avalon;
using namespace Avalon::IO;
using namespace Avalon::Network;
using namespace Avalon::Threading;
using namespace boost;
using namespace boost::asio;
using namespace boost::asio::ip;
using namespace boost::posix_time;

namespace {
  const int DEFAULT_READ_SIZE = 1024;
}

struct MulticastSocket::SocketDetails {
  boost::asio::ip::address_v4 m_group;
  boost::asio::ip::address_v4 m_interface;
  boost::asio::ip::udp::endpoint m_listenEndpoint;
  boost::asio::ip::udp::endpoint m_senderEndpoint;
  boost::scoped_ptr<boost::asio::ip::udp::socket> m_socket; 

  SocketDetails(const IpAddress& group, const IpAddress& interface)
      : m_group(ip::address_v4::from_string(group.first)),
        m_interface(ip::address_v4::from_string(interface.first)), 
        m_listenEndpoint(ip::udp::v4(), group.second) {}
};

MulticastSocket::MulticastSocket(const IpAddress& group,
    const IpAddress& interface)
    : m_socketDetails(new SocketDetails(group, interface)),
      m_isOpen(false),
      m_receiveBufferSize(0),
      m_pendingClose(false),
      m_receiveSlot(bind(&MulticastSocket::HandleReceive, this, _1, _2)),
      m_sendSlot(bind(&MulticastSocket::HandleSend, this, _1, _2)) {}

MulticastSocket::~MulticastSocket() {
  Close();
}

int MulticastSocket::GetReceiveBufferSize() const {
  return m_receiveBufferSize;
}

void MulticastSocket::SetReceiveBufferSize(int size) {
  assert(!m_isOpen);
  m_receiveBufferSize = size;
}

void MulticastSocket::Open() {
  lock_guard<mutex> lock(m_mutex);
  if(m_isOpen) {
    return;
  }
  m_socketDetails->m_socket.reset(new ip::udp::socket(
    SocketThreadPool::GetInstance().GetService(),
    m_socketDetails->m_listenEndpoint.protocol()));
  m_socketDetails->m_socket->set_option(ip::udp::socket::reuse_address(true));
  m_socketDetails->m_socket->bind(m_socketDetails->m_listenEndpoint);
  if(m_receiveBufferSize != 0) {
    socket_base::receive_buffer_size bufferSize(m_receiveBufferSize);
    m_socketDetails->m_socket->set_option(bufferSize);
  }
  m_socketDetails->m_socket->set_option(ip::multicast::join_group(
    m_socketDetails->m_group, m_socketDetails->m_interface));
  m_isOpen = true;
}

void MulticastSocket::Close() {
  unique_lock<mutex> lock(m_mutex);
  if(!m_isOpen) {
    return;
  }
  m_socketDetails->m_socket->close();
  m_pendingClose = true;
  while(m_receiveParameters.m_packet != NULL ||
      m_sendParameters.m_result != NULL) {
    m_closedCondition.wait(lock);
  }
  m_pendingClose = false;
  m_isOpen = false;
}

void MulticastSocket::Receive(Buffer* buffer,
    Out<Async<DatagramPacket> > packet) {
  lock_guard<mutex> lock(m_mutex);
  m_receiveParameters.m_buffer = buffer;
  m_receiveParameters.m_bufferOffset = buffer->GetSize();
  m_receiveParameters.m_packet = packet.Get();
  buffer->Grow(DEFAULT_READ_SIZE);
  m_socketDetails->m_socket->async_receive_from(boost::asio::buffer(
    buffer->GetMutableData() + m_receiveParameters.m_bufferOffset,
    DEFAULT_READ_SIZE), m_socketDetails->m_senderEndpoint, m_receiveSlot);
}

void MulticastSocket::Send(const DatagramPacket& packet,
    Out<Async<void> > result) {
  lock_guard<mutex> lock(m_mutex);
  Buffer dataCopy = packet.GetData();
  ip::udp::endpoint destination(ip::address::from_string(
    packet.GetAddress().first), packet.GetAddress().second);
  m_socketDetails->m_socket->async_send_to(buffer(dataCopy.GetData(),
    dataCopy.GetSize()), destination, m_sendSlot);
}

void MulticastSocket::HandleReceive(const boost::system::error_code& error,
    size_t size) {
  unique_lock<mutex> lock(m_mutex);
  if(error != 0) {
    m_receiveParameters.m_buffer->Shrink(DEFAULT_READ_SIZE);
    Async<DatagramPacket>* packet = m_receiveParameters.m_packet;
    m_receiveParameters.m_packet = NULL;
    if(m_pendingClose) {
      m_closedCondition.notify_one();
    }
    lock.unlock();
    packet->SetException(SocketException(error.value(), error.message()));
    return;
  }
  m_receiveParameters.m_buffer->Shrink(DEFAULT_READ_SIZE - size);
  Async<DatagramPacket>* result = m_receiveParameters.m_packet;
  m_receiveParameters.m_packet = NULL;
  DatagramPacket packet(static_cast<int>(size), *m_receiveParameters.m_buffer,
    make_pair(m_socketDetails->m_senderEndpoint.address().to_string(),
    m_socketDetails->m_senderEndpoint.port()));
  if(m_pendingClose) {
    m_closedCondition.notify_one();
  }
  lock.unlock();
  result->SetResult(packet);
}

void MulticastSocket::HandleSend(const boost::system::error_code& error,
    size_t size) {
  unique_lock<mutex> lock(m_mutex);
  if(error != 0) {
    Async<void>* result = m_sendParameters.m_result;
    m_sendParameters.m_result = NULL;
    if(m_pendingClose) {
      m_closedCondition.notify_one();
    }
    lock.unlock();
    result->SetException(SocketException(error.value(), error.message()));
    return;
  }
  Async<void>* result = m_sendParameters.m_result;
  m_sendParameters.m_result = NULL;
  if(m_pendingClose) {
    m_closedCondition.notify_one();
  }
  lock.unlock();
  result->SetResult();
}
